from langchain.document_loaders import TextLoader,JSONLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
import sys
from faissManager import FAISS
from sent2vec import Sent2VecEmbeddings
from langchain.document_transformers import (
    LongContextReorder,
)

from dataclasses import dataclass

import numpy as np
# @dataclasss
# class template:

class MultiTurnRetriever:
    def __init__(self,index_dir='/data/lxy/RAT/10m9d_wxb_bge1.5_hnswivf_20500',model_name='BAAI/bge-large-zh-v1.5',
                 only_retriever=False,nsp_classifier_model=None,nsp_classifier_tokenizer=None,isEnsemble=False,model=None) -> None:
        self.only_retriever=only_retriever
        if model is None:
            # 考虑到有可能在外面初始化。
            model=Sent2VecEmbeddings(model_name=model_name)
        self.db=FAISS.load_local(index_dir,model)
        self.isEnsemble=isEnsemble
        if not isEnsemble:
            self.nsp_classifier_tokenizer=nsp_classifier_tokenizer
            self.nsp_classifier_model=nsp_classifier_model
            self.nsp_classifier_model.cuda()
        
        print(f'成功加载了编码器{model_name},和向量数据库{index_dir}-------------------------------------------------------------------------------------------------------')
        print('==========================================================================================================================================================')

    def fuseHistoryEmbeddings(self,historyQueryEmbeddingList,debug):
        if debug:
            print('调试，历史查询变量列表',historyQueryEmbeddingList)
        if len(historyQueryEmbeddingList):
            # 目前先假定传进来的都是同一族的吧
            historyQueryLen=len(historyQueryEmbeddingList)
            # 按照历史分布权重，比如有三轮，那么就是 1/1+2+3, 2/ 3/... ,从而使得上一轮的问题会得到更多的重视
            weights_for_sum=[x/sum(range(1,historyQueryLen+1)) for x in range(1,historyQueryLen+1)]
            return np.average(historyQueryEmbeddingList,weights=weights_for_sum,axis=0)
        else:
            return None
        
    def isSameType(self,query,history,debug):
        # 其实本来可以送history[-1]进来，但是如果是空会越界，在外面判断不合适
        if len(history)==0:
            return True
        if debug:
            print('调试',history[-1],query)
        inputs=self.nsp_classifier_tokenizer(history[-1],query,return_tensors='pt',truncation=True,
            max_length=512).to('cuda')
        ans=self.nsp_classifier_model(**inputs).logits

        if ans[0,0]>ans[0,1]:
            return False
        return True
    def retrieve(self,query,k=3,score_threshold=0.4,roles=None,debug=False,
                 history_embedding=[],border=[0,0],return_list=False):
        # 把第二行参数放在形参而不是成员变量，是考虑到不同用户发起的不同访问信息会在这里累计，因此需要在gradio里维护。
        # history是一个交替的query和response列表，我们似乎只需要用到最后一个response.如果直接使用外部的history，这里我们就不需要维护了。
        # history_embedding 历史所有query的embedding缓存
        # border是用来划分history_embedding里的相同语义块,每两个元素之间会形成一个左闭右开区间
        if debug:
            print('调试，query:',query)
        # 先判断正确的question是什么
        if roles is not None:
                # 说明是fastchat，传进来的应该是[role1,role2]
                real_question=query[query.rfind(roles[0])+len(roles[0]):query.rfind(roles[1])]
        else:
            real_question=query

        if self.isEnsemble:
            # 退化的单轮搜索，要注意这个可能因为阈值返回空
            docs_single,scores_single,sentence_embedding = self.db.similarity_search(real_question,10*k,score_threshold=score_threshold,
                                                            multiR=False,history_embedding=None)
            
            # 先假定当前轮应该做融合,且和历史块是同义块
            border[-1]+=1
            historyQueryEmbeddingList=history_embedding[border[-2]:border[-1]]
            history_query_embedding=self.fuseHistoryEmbeddings(historyQueryEmbeddingList,debug=debug)
            if debug:
                print('调试，在融合前的同义块划分',border)
            docs_mult,scores_mult,sentence_embedding = self.db.similarity_search(real_question,10*k,score_threshold=score_threshold,
                                                            multiR=True,history_embedding=history_query_embedding)
            if debug:
                print('调试，分数对比',f'单轮分数:{scores_single[:k]},{sum(scores_single[:k])}','多轮分数：',scores_mult[:k],sum(scores_mult[:k]))
            

            docs=docs_single
            # 这里可能有个问题，就是单轮和多轮的搜索如果设置阈值，可能会返回空，使得结果不真实，因为列表是空就sum就是0
            # 但如果不设置阈值，可能要重新搜一次。\
            if len(scores_single)<len(scores_mult):
                if debug:
                    print('调试','该问题单轮召回数目少于多轮，认定为不同块，但是仍然采用多轮的检索结果（防止单轮是空的情况）')
                # 先还原上面的假定
                border[-1]-=1
                border.append(border[-1]+1)
                docs=docs_mult
            elif sum(scores_single[:k])<sum(scores_mult[:k]): # 说明融合距离反而变远了，认定为非同义块
                if debug:
                    print('调试','a非同一块')
                # 先还原上面的假定
                border[-1]-=1
                border.append(border[-1]+1)
            else:
                if debug:
                    print('调试','a是同一块')
                # 这里不需要做别的事情，因为上面假定为同义块,只需要把docs选择为多轮块就好了
                docs=docs_mult
            if debug:
                print('调试，融合后的同义块划分',border)
            history_embedding.append(sentence_embedding)

            # 一个大胆的尝试，当前轮的回答和同义块划分应当没有关系，直接把单多轮的结果混起来按照分数返回。
            
            total=list(zip(docs_mult,scores_mult,[1 for _ in range(len(docs_mult))]))+list(zip(docs_single,scores_single,[0 for _ in range(len(docs_single))]))

            # 按分数先排序一下，但是很可能有重复的，得筛！
            total=sorted(total,key=lambda x:x[1])
            docs=total
            # 5. 滤除检索结果里可能的重复项
            docs_content=set()
            valid_doc_count=0
            valid_doc=[]
            debug_use_for_source=[]
            for i,doc in enumerate(docs):
                if valid_doc_count>=k:
                    break
                if doc[0].page_content not in docs_content:
                    docs_content.add(doc[0].page_content)
                    valid_doc_count+=1
                    valid_doc.append(doc[0])
                    debug_use_for_source.append(doc[2])
            if debug:
                print('混合检索，目前是不稳定特性。单轮来源是0，多轮来源是1',debug_use_for_source)
        else:
            # 1.需要先判断当前输入和当前族是否同属一族
            if self.isSameType(query,history,debug=debug):
                if debug:
                    print('调试','是同一块')
                border[-1]+=1
                # 2. 然后进行历史信息内部的融合
                historyQueryEmbeddingList=history_embedding[border[-2]:border[-1]]
                history_query_embedding=self.fuseHistoryEmbeddings(historyQueryEmbeddingList)
            else:
                if debug:
                    print('调试','非同一块')
                border.append(border[-1]+1)
                history_query_embedding=None
            if debug:
                print('调试',border)
            
                
            if debug:
                # print(f'获取到的问题：{real_question}\n')
                print('马上要融合历史信息和当前信息了，历史信息是',history_query_embedding)
            
            # 3. 融合历史和当前信息，得到当前检索结果
            docs,scores,sentence_embedding = self.db.similarity_search(real_question,10*k,score_threshold=score_threshold,
                                                                multiR=True,history_embedding=history_query_embedding)
                
            # 4. 将当前的embedding保存到缓存中，不要保存query，因为history在外面是自动维护的。
            history_embedding.append(sentence_embedding)
        
        
            # 5. 滤除检索结果里可能的重复项
            docs_content=set()
            valid_doc_count=0
            valid_doc=[]
            for i,doc in enumerate(docs):
                if valid_doc_count>=k:
                    break
                if doc.page_content not in docs_content:
                    docs_content.add(doc.page_content)
                    valid_doc_count+=1
                    valid_doc.append(doc)
            


        # 6.做重排序，但是在仅检索的时候不需要重排序，因为reorder只和大模型特性相关。
        reordering_docs=valid_doc
        if self.only_retriever==False:
            reordering = LongContextReorder()
            reordering_docs = reordering.transform_documents(valid_doc)
        if debug:
            print(f'过滤掉了{10*k-len(docs)}个距离超出{score_threshold}检索到的文档，经过去重又过滤了{len(docs)-len(valid_doc)}个\n')
        
        # 7.确定返回模式，只检索的话，带着检索结果回去就行，如果是生成模式，还需要拼接一定的prompt
        if self.only_retriever==True:
            if len(reordering_docs)==0:
                return 0
            answer=''
            for i in range(len(reordering_docs)):
                answer+=reordering_docs[i].page_content
            return answer
        if return_list:
            return [d.page_content for d in reordering_docs]
        if len(reordering_docs):
            prompt="背景：\n"
            # prompt=""
            for i,d in enumerate(reordering_docs):
                prompt+=f'[{i+1}].'+reordering_docs[i].page_content+'\n'

            # prompt=query[:begin_idx]+prompt+'问题：'+real_question+query[end_idx:]
            
            if roles is None:
                prompt=prompt+"\n\n问题: "+real_question
                return prompt
            else:
                # 是fastchat！
                    #  sys prompt + knowledge prompt + role1 +real_q +role2
                return query[:query.rfind(roles[0])]+prompt+query[query.rfind(roles[0]):query.rfind(roles[0])+len(roles[0])]+real_question+query[query.rfind(roles[1]):],sentence_embedding
        else:
            return query
        
    def parseHistory(self,history,debug=False):
        # 现在目标是，从一个历史里算出retrieve所需要的history_embedding和border
        history_embedding=[]
        border=[0,0]
        if debug:
            print('调试,编码历史=============================================================')
            print('history',history)
        for i in range(0,len(history),2):
            self.retrieve(history[i],k=3,score_threshold=0.4,roles=None,debug=debug,
                 history_embedding=history_embedding,border=border,return_list=True)
        if debug:
            print('调试,编码历史',history_embedding,border)
            print('调试,编码历史=============================================================')
        return history_embedding,border
        # return self.retrieve(query,k=3,score_threshold=0.4,roles=None,debug=True,
        #          history=[],history_embedding=[],border=[0,0],return_list=True,isEnsemble=False)

                
        
if __name__=='__main__':
    index_dir='/home/lxy/multiR/gte_11m21d'
    model_name='/data/lxy/thenlper/gte-large-zh'
    # model_name='BAAI/bge-large-zh-v1.5'
    
    r=MultiTurnRetriever(index_dir,model_name=model_name,isEnsemble=True)
    history=['怎么申请图书馆报告厅','我也不知道啊','南湖的','哦','你是谁？','不告诉你','再见']
    # print(r.retrieve("我想办理在校生四六级证明"))
    import time
    start_time = time.time()
    r.parseHistory(history)
    print("--- %s seconds ---" % (time.time() - start_time))

